import os
import shutil
from PIL import Image
import cv2
import pytesseract
# 设置Tesseract的路径（如果它不在默认路径中）
# pytesseract.pytesseract.tesseract_cmd = r'E:\software\devtools\Tesseract-OCR\tesseract.exe'  # 你的Tesseract安装路径

# 设置Tesseract的配置参数，限制识别的字符集为数字
custom_config = r'--psm 6 outputbase digits'

rules = {
    'SLC SP Erase': '11-',
    'SLC SP PGM': '12-',
    'SLC SP Read': '13-',
    'SLC MP Erase': '21-',
    'SLC MP PGM': '22-',
    'SLC MP Read': '23-',
    'TLC SP Erase': '31-',
    'TLC SP PGM': '32-',
    'TLC SP Read-LP': '33-',
    'TLC SP Read-MP': '34-',
    'TLC SP Read-UP': '35-',
    'TLC MP Erase': '41-',
    'TLC MP PGM': '42-',
    'TLC MP Read-LP': '43-',
    'TLC MP Read-MP': '44-',
    'TLC MP Read-UP': '45-'
}

def modify_filenames(folder_path, rules):
    # 遍历文件夹中的文件名
    for filename in os.listdir(folder_path):
        # 获取文件的完整路径
        file_path = os.path.join(folder_path, filename)

        # 检查文件名是否已经包含了前缀
        has_prefix = False
        for prefix in rules.values():
            if filename.startswith(prefix):
                has_prefix = True
                break

        # 如果文件名已经包含了前缀，跳过重命名的步骤
        if has_prefix:
            continue

        # 遍历规则
        for text, prefix in rules.items():
            # 如果文件名中包含规则中的文本
            if text in filename:
                # 构建新的文件名
                new_filename = prefix + filename

                # 重命名文件
                os.rename(file_path, os.path.join(folder_path, new_filename))

                # 打印修改后的文件名
                # print(new_filename)

                # 跳出内层循环，继续处理下一个文件名
                break


def process_folder(folder_path):
    # 获取文件夹下的所有文件名
    file_names = os.listdir(folder_path)
    
    # 创建一个字典用于存储前缀和对应的文件名列表
    prefix_dict = {}
    
    # 遍历文件夹下的每个文件名
    for file_name in file_names:
        # 将文件名按照"-"分割成前缀和文件名的列表
        parts = file_name.split("-")
        prefix = parts[0]
        
        # 如果前缀已经存在于prefix_dict中，则将文件名添加到对应的列表中
        # 否则，创建一个新的列表并将文件名添加进去
        if prefix in prefix_dict:
            prefix_dict[prefix].append(file_name)
        else:
            prefix_dict[prefix] = [file_name]
    
    # 创建一个新的文件夹用于存放多余的文件
    new_folder_path = os.path.join(folder_path, "else")
    os.makedirs(new_folder_path, exist_ok=True)
    
    # 遍历prefix_dict中的每个前缀和文件名列表
    for prefix, file_list in prefix_dict.items():
        # 如果文件名列表的长度大于1，将多余的文件移动到新的文件夹下
        if len(file_list) > 1:
            for file_name in file_list[1:]:
                old_file_path = os.path.join(folder_path, file_name)
                new_file_path = os.path.join(new_folder_path, file_name)
                shutil.move(old_file_path, new_file_path)
    
    # 返回处理后的文件夹下的文件名列表
    return os.listdir(folder_path)


def ensure_directory_exists(directory):
    """
    Ensure that the directory exists, create it if necessary.
    
    Args:
    - directory (str): Path to the directory.
    """
    os.makedirs(directory, exist_ok=True)

def clear_folder(folder_path):
    # 检查文件夹是否存在
    if not os.path.exists(folder_path):
        print(f"文件夹 {folder_path} 不存在，无需清空")
        return

    # 提示用户是否确定清空
    # choice = input(f"确定要清空文件夹 {folder_path} 吗？(Y/N): ")
    # if choice.lower() != 'y':
    #     print("取消清空")
    #     return

    # 清空文件夹
    shutil.rmtree(folder_path)

    print(f"文件夹 {folder_path} 已清空")

def preprocess_image(file_path, output_path, threshold=255):
    # 读取图像
    image = cv2.imread(file_path)

    # 将图像转换为灰度图
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    # 对图像进行二值化处理
    _, binary = cv2.threshold(gray, 0, 256, cv2.THRESH_BINARY_INV | cv2.THRESH_OTSU)

    cv2.imwrite(output_path, binary)

def preprocess_image_with_color(image_path, output_path, threshold_rgb):
    threshold_r = threshold_rgb[0]
    threshold_g = threshold_rgb[1]
    threshold_b = threshold_rgb[2]
    # 打开图像
    image = Image.open(image_path)
    
    # 创建一个新的图像，将深色保留，浅色去除
    processed_image = Image.new('RGB', image.size)
    for x in range(image.width):
        for y in range(image.height):
            pixel = image.getpixel((x, y))
            if (pixel[0] < threshold_r and pixel[1] < threshold_g and pixel[2] < threshold_b):
                processed_image.putpixel((x, y), (255, 255, 255)) # 白色，不保留
            else:
                processed_image.putpixel((x, y), (0, 0, 0))
    
    # 保存处理后的图像
    processed_image.save(output_path)

def crop_image(file_path, output_path, crop_pos):
    """
    Crop the image and save it to the output path.
    
    Args:
    - file_path (str): Path to the input image file.
    - output_path (str): Path to save the cropped image.
    """
    # Specify the coordinates for cropping (left, top, right, bottom)
    left = crop_pos["left"] # x-coordinate of the left edge
    top = crop_pos["top"]   # y-coordinate of the upper edge
    right = crop_pos["right"] # x-coordinate of the right edge
    bottom = crop_pos["bottom"] # y-coordinate of the lower edge
    
    img = Image.open(file_path)
    cropped_img = img.crop((left, top, right, bottom))
    cropped_img.save(output_path)


def recognize_digits(file_path, output_path):
    # 使用Tesseract进行数字识别
    # config = "--psm 6"
    # digits = pytesseract.image_to_string(file_path, config=config)
    digits = pytesseract.image_to_string(file_path)

    with open(output_path, 'w', encoding='utf-8') as f:
        f.write(digits)

# def process_images_in_folder(folder_path, output_folder, process_func, output_extension):
#     """
#     Process all PNG images in the given folder.
    
#     Args:
#     - folder_path (str): Path to the folder to be processed.
#     - output_folder (str): Path to the folder where the processed images will be saved.
#     - process_func (function): The function used to process the image.
#     - output_extension (str): Extension of the output files.
#     """
#     clear_folder(output_folder)
#     for root, _, files in os.walk(folder_path):
#         for file in files:
#             if file.endswith('.png'):
#                 file_path = os.path.join(root, file)
#                 output_path = os.path.join(output_folder, os.path.splitext(file)[0] + output_extension)
#                 os.makedirs(os.path.dirname(output_path), exist_ok=True)
#                 process_func(file_path, output_path)
def process_images_in_folder(folder_path, output_folder, process_func, output_extension, *args, **kwargs):
    """
    Process all PNG images in the given folder.
    
    Args:
    - folder_path (str): Path to the folder to be processed.
    - output_folder (str): Path to the folder where the processed images will be saved.
    - process_func (function): The function used to process the image.
    - output_extension (str): Extension of the output files.
    - *args: Additional positional arguments to be passed to the process_func.
    - **kwargs: Additional keyword arguments to be passed to the process_func.
    """
    clear_folder(output_folder)
    for root, _, files in os.walk(folder_path):
        for file in files:
            if file.endswith('.png'):
                file_path = os.path.join(root, file)
                output_path = os.path.join(output_folder, os.path.splitext(file)[0] + output_extension)
                os.makedirs(os.path.dirname(output_path), exist_ok=True)
                process_func(file_path, output_path, *args, **kwargs)

def process_images_in_folder_crop(folder_path, output_folder, crop_pos):
    process_images_in_folder(folder_path, output_folder, crop_image, '.png', crop_pos)

def process_images_in_folder_preprocess_with_color(folder_path, output_folder, threshold_rgb):
    process_images_in_folder(folder_path, output_folder, preprocess_image_with_color, '.png', threshold_rgb)

def process_images_in_folder_preprocess(folder_path, output_folder, threshold):
    process_images_in_folder(folder_path, output_folder, preprocess_image, '.png', threshold)

def process_images_in_folder_ocr(folder_path, output_folder):
    process_images_in_folder(folder_path, output_folder, recognize_digits, '.txt')

def join_path(path, folder):
    return path + folder

def process_images(data_path, output_path, crop_func, crop_pos, preprocess_func, threshold):
    crop_path = join_path(output_path, "1_crop")
    crop_func(data_path, crop_path, crop_pos)
    modify_filenames(crop_path, rules)
    process_folder(crop_path)

    preprocess_path = join_path(output_path, "2_preporcess")
    preprocess_func(crop_path, preprocess_path, threshold)
    process_folder(preprocess_path)


    output_folder = join_path(output_path, "3_text")
    process_images_in_folder_ocr(preprocess_path, output_folder)
    process_folder(output_folder)


